# Created on 2018/12
# Author: Kaituo XU
import math
import mindspore
import mindspore.ops as ops
from mindspore import Tensor
import numpy as np
from mindspore import context


def overlap_and_add(signal, frame_step):
    """Reconstructs a signal from a framed representation.

    Adds potentially overlapping frames of a signal with shape
    `[..., frames, frame_length]`, offsetting subsequent frames by `frame_step`.
    The resulting tensor has shape `[..., output_size]` where

        output_size = (frames - 1) * frame_step + frame_length

    Args:
        signal: A [..., frames, frame_length] Tensor. All dimensions may be unknown, and rank must be at least 2.
        frame_step: An integer denoting overlap offsets. Must be less than or equal to frame_length.

    Returns:
        A Tensor with shape [..., output_size] containing the overlap-added frames of signal's inner-most two dimensions.
        output_size = (frames - 1) * frame_step + frame_length

    Based on https://github.com/tensorflow/tensorflow/blob/r1.12/tensorflow/contrib/signal/python/ops/reconstruction_ops.py
    """
    outer_dimensions = signal.shape[:-2]
    frames, frame_length = signal.shape[-2:]

    subframe_length = math.gcd(frame_length, frame_step)  # gcd=Greatest Common Divisor
    subframe_step = frame_step // subframe_length
    subframes_per_frame = frame_length // subframe_length
    output_size = frame_step * (frames - 1) + frame_length
    output_subframes = output_size // subframe_length

    subframe_signal = signal.view(*outer_dimensions, -1, subframe_length)

    #frame = mindspore.numpy.arange(0, output_subframes).unfold(0, subframes_per_frame, subframe_step)
    frame = mindspore.numpy.arange(0, output_subframes)
    frame = ops.Concat(-1)((ops.expand_dims(frame[0:-1:subframe_step], 1), ops.expand_dims(frame[1::subframe_step], 1)))
    # frame = signal.new_tensor(frame).long()  # signal may in GPU or CPU
    frame = frame.view(-1)

    # zeros = ops.Zeros()
    # # [2, 2, 32000,1]
    # result = zeros((*outer_dimensions, output_subframes, subframe_length), mindspore.float32)
    zeros = ops.Zeros()
    # [2, 2, 32000,1]
    # result = zeros((*outer_dimensions, output_subframes, subframe_length), mindspore.float32)
    pad = zeros((subframe_signal.shape[0], subframe_signal.shape[1], 1, subframe_signal.shape[3]), mindspore.float32)
    conc = ops.Concat(2)
    subframe_signal = conc((pad, subframe_signal, pad))
    subframe_signal_first = subframe_signal[:, :, ::2, :]
    subframe_signal_right = subframe_signal[:, :, 1::2, :]
    result = subframe_signal_first + subframe_signal_right
    # i = 0
    # j = 0
    # while (j < subframe_signal.shape[2]) and (i < len(frame)):
    #     if i == (len(frame) - 1) or frame[i] != frame[i + 1]:
    #         result[:, :, frame[i], :] += subframe_signal[:, :, j, :]
    #         j = j + 1
    #         i = i + 1
    #     else:
    #         result[:, :, frame[i], :] += subframe_signal[:, :, j, :]
    #         result[:, :, frame[i], :] += subframe_signal[:, :, j + 1, :]
    #         j = j + 2
    #         i = i + 2
    # transpose = ops.Transpose()
    # result = transpose(result, (2, 1, 0, 3))
    #
    # frame = torch.unsqueeze(frame, 1).repeat(1, subframe_signal)
    # indice = frame
    # result = ops.ScatterNdAdd(result, indice, subframe_signal)
    # # indices = frame
    # # inplaceAdd = ops.InplaceAdd(indices)
    # # result = inplaceAdd(result, subframe_signal)
    # result = transpose(result, (2, 1, 0, 3))
    result = result.view(*outer_dimensions, -1)
    return result


def remove_pad(inputs, inputs_lengths):
    """
    Args:
        inputs: torch.Tensor, [B, C, T] or [B, T], B is batch size
        inputs_lengths: torch.Tensor, [B]
    Returns:
        results: a list containing B items, each item is [C, T], T varies
    """
    results = []
    dim = inputs.dim()
    if dim == 3:
        C = inputs.size(1)
    for input, length in zip(inputs, inputs_lengths):
        if dim == 3: # [B, C, T]
            results.append(input[:,:length].view(C, -1).asnumpy())
        elif dim == 2:  # [B, T]
            results.append(input[:length].view(-1).asnumpy())
    return results

if __name__ == "__main__":
    context.set_context(mode=context.PYNATIVE_MODE, device_target="CPU")
    signal = np.random.randn(2, 2, 31999, 2).astype(np.float32)
    signal =Tensor(signal)
    frame_step = 1
    result = overlap_and_add(signal, frame_step)